#!/usr/bin/env python3
"""
Input File Gantt Chart Visualization
Visualizes memory access patterns and I/O requirements from input files
"""

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches
from matplotlib.patches import Rectangle
import numpy as np
from pathlib import Path

def parse_input_file(filepath):
    """Parse input file and extract request information"""
    with open(filepath, 'r') as f:
        first_line = f.readline().split()
        L = int(first_line[0])  # Virtual address space
        M = int(first_line[1])  # HBM capacity
        N = int(first_line[2])  # Number of requests
        
        requests = []
        for i in range(N):
            line = f.readline().split()
            requests.append({
                'id': i,
                'addr': int(line[0]),
                'size': int(line[1]),
                'start': int(line[2]),
                'duration': int(line[3])
            })
    
    return L, M, N, requests


def visualize_input_gantt(filepath, output_path):
    """
    Create a Gantt chart showing:
    1. Request computation time
    2. Required data loading time (theoretical)
    3. Memory access patterns
    """
    
    # Parse input
    L, M, N, requests = parse_input_file(filepath)
    
    # Calculate metrics
    IO_CYCLES_PER_BYTE = 40
    
    # Create figure with 3 subplots (increased spacing)
    fig, (ax_compute, ax_load, ax_memory) = plt.subplots(
        3, 1, 
        figsize=(16, 12),
        gridspec_kw={'height_ratios': [2, 2, 1.5], 'hspace': 0.4}
    )
    
    # Color scheme
    colors = plt.cm.Set3(np.linspace(0, 1, N))
    
    # ==========================================
    # Subplot 1: Computation Time (NPU)
    # ==========================================
    ax_compute.set_title(
        f'Visit Time\n'
        f'L: {L:,} | M: {M:,} | N: {N}',
        fontsize=14, weight='bold', pad=20
    )
    ax_compute.set_ylabel('Request', fontsize=12, weight='bold')
    ax_compute.set_xlabel('Time', fontsize=12, weight='bold')
    
    for req in requests:
        ax_compute.barh(
            y=req['id'],
            width=req['duration'],
            left=req['start'],
            height=0.6,
            color=colors[req['id']],
            edgecolor='black',
            linewidth=1.5,
            alpha=0.8,
            label=f"R{req['id']}"
        )
    
    ax_compute.set_yticks(range(N))
    ax_compute.set_yticklabels([f'R{i}' for i in range(N)])
    ax_compute.invert_yaxis()
    ax_compute.grid(axis='x', linestyle=':', alpha=0.6)
    ax_compute.set_xlim(left=0, right=max(r['start'] + r['duration'] for r in requests) * 1.1)
    
    # ==========================================
    # Subplot 2: Required I/O Time (Theoretical)
    # ==========================================
    ax_load.set_title(
        'I/O Time',
        fontsize=14, weight='bold', pad=15
    )
    ax_load.set_ylabel('Request', fontsize=12, weight='bold')
    ax_load.set_xlabel('Time', fontsize=12, weight='bold')
    
    for req in requests:
        load_time = req['size'] * IO_CYCLES_PER_BYTE
        
        ax_load.barh(
            y=req['id'],
            width=load_time,
            left=req['start'],
            height=0.6,
            color=colors[req['id']],
            edgecolor='darkred',
            linewidth=2,
            alpha=0.7,
            hatch='//'
        )
    
    ax_load.set_yticks(range(N))
    ax_load.set_yticklabels([f'R{i}' for i in range(N)])
    ax_load.invert_yaxis()
    ax_load.grid(axis='x', linestyle=':', alpha=0.6)
    
    # Calculate max time considering I/O
    max_io_time = max(r['start'] + r['size'] * IO_CYCLES_PER_BYTE for r in requests)
    ax_load.set_xlim(left=0, right=max_io_time * 1.1)
    
    # ==========================================
    # Subplot 3: Memory Address Space Usage
    # ==========================================
    ax_memory.set_title(
        f'L: {L:,} | M: {M:,} | N: {N}',
        fontsize=14, weight='bold', pad=15
    )
    ax_memory.set_ylabel('Request', fontsize=12, weight='bold')
    ax_memory.set_xlabel('Memory Address', fontsize=12, weight='bold')
    
    for req in requests:
        ax_memory.barh(
            y=req['id'],
            width=req['size'],
            left=req['addr'],
            height=0.6,
            color=colors[req['id']],
            edgecolor='black',
            linewidth=1.5,
            alpha=0.6
        )
    
    # Add HBM capacity line
    ax_memory.axvline(x=M, color='red', linestyle='--', linewidth=2, label=f'HBM Capacity ({M:,})')
    ax_memory.axvline(x=L, color='blue', linestyle='--', linewidth=2, label=f'Address Space ({L:,})')
    
    ax_memory.set_yticks(range(N))
    ax_memory.set_yticklabels([f'R{i}' for i in range(N)])
    ax_memory.invert_yaxis()
    ax_memory.grid(axis='x', linestyle=':', alpha=0.6)
    ax_memory.set_xlim(left=0, right=L * 1.05)
    ax_memory.legend(loc='upper right', fontsize=10)
    
    # ==========================================
    # Summary Statistics (Text Box)
    # ==========================================
    
    # Calculate statistics
    total_compute_time = sum(r['duration'] for r in requests)
    total_data_size = sum(r['size'] for r in requests)
    unique_bytes = set()
    for req in requests:
        for addr in range(req['addr'], req['addr'] + req['size']):
            unique_bytes.add(addr)
    unique_data_size = len(unique_bytes)
    theoretical_io_time = unique_data_size * IO_CYCLES_PER_BYTE
    
    io_density = unique_data_size / total_compute_time if total_compute_time > 0 else 0
    memory_pressure = unique_data_size / M
    time_pressure = theoretical_io_time / total_compute_time if total_compute_time > 0 else 0
    
    # stats_text = (
    #     f"═══ Workload Statistics ═══\n"
    #     f"Total Compute Time: {total_compute_time:,} cycles\n"
    #     f"Total Data Requested: {total_data_size:,} bytes\n"
    #     f"Unique Data: {unique_data_size:,} bytes\n"
    #     f"Theoretical I/O Time: {theoretical_io_time:,} cycles\n"
    #     f"\n"
    #     f"I/O Density Ratio: {io_density / 0.025:.2f}x\n"
    #     f"Memory Pressure: {memory_pressure:.2f}x (HBM)\n"
    #     f"Time Pressure: {time_pressure:.1%}\n"
    #     f"\n"
    #     f"Workload Type: "
    #     f"{'IO-Intensive' if io_density / 0.025 > 1.05 else 'Compute-Intensive'} + "
    #     f"{'Memory-Intensive' if memory_pressure >= 0.95 else 'Memory-Sufficient'}"
    # )
    
    # fig.text(
    #     0.02, 0.02, stats_text,
    #     fontsize=10,
    #     family='monospace',
    #     bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5),
    #     verticalalignment='bottom'
    # )
    
    # ==========================================
    # Layout and Save
    # ==========================================
    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches='tight')
    plt.close()
    
    print(f"✓ Gantt chart saved to: {output_path}")



def main():
    """Main function to process input files"""
    
    # Default paths
    input_file = Path("/Users/chenxi/Documents/比赛/CCF/checker/example/infile/infile_32.txt")
    output_file = Path("/Users/chenxi/Documents/比赛/CCF/checker/example/fig/gantt_input_32.png")
    
    # Ensure output directory exists
    output_file.parent.mkdir(parents=True, exist_ok=True)
    
    print("=" * 60)
    print("Input File Gantt Chart Visualization")
    print("=" * 60)
    print(f"Input:  {input_file}")
    print(f"Output: {output_file}")
    print("-" * 60)
    
    if not input_file.exists():
        print(f"Error: Input file not found: {input_file}")
        return
    
    visualize_input_gantt(input_file, output_file)
    
    print("=" * 60)
    print("Visualization complete!")
    print("=" * 60)


if __name__ == "__main__":
    main()
